"""
KSQL相关API路由
基于Confluent KSQL REST API官方文档实现
"""
import json
import asyncio
from typing import Optional, Dict, Any, List
from fastapi import APIRouter, HTTPException, BackgroundTasks
from pydantic import BaseModel
import aiohttp
import base64

from app.core.redis_client import redis_manager
from app.utils.logger import get_logger

router = APIRouter()
logger = get_logger(__name__)


class KsqlQueryRequest(BaseModel):
    """KSQL查询请求"""
    ksql: str
    streamsProperties: Optional[Dict[str, Any]] = {}
    sessionVariables: Optional[Dict[str, Any]] = {}


class KsqlQueryResponse(BaseModel):
    """KSQL查询响应"""
    success: bool
    data: Optional[List[Dict[str, Any]]] = None
    error: Optional[str] = None
    execution_time: Optional[int] = None
    row_count: Optional[int] = None
    query_id: Optional[str] = None
    command_status: Optional[Dict[str, Any]] = None


class KsqlStreamRequest(BaseModel):
    """KSQL流查询请求"""
    sql: str
    properties: Optional[Dict[str, Any]] = {}


async def get_ksql_config():
    """获取KSQL配置"""
    try:
        config_data = await redis_manager.get_json("ksql:config:connection")
        if not config_data or not config_data.get('ksql_enabled'):
            raise HTTPException(status_code=400, detail="KSQL未启用或未配置")
        return config_data
    except Exception as e:
        logger.error(f"获取KSQL配置失败: {e}")
        raise HTTPException(status_code=500, detail="获取KSQL配置失败")


def build_headers(config: Dict[str, Any]) -> Dict[str, str]:
    """构建请求头"""
    headers = {
        'Content-Type': 'application/vnd.ksql.v1+json; charset=utf-8',
        'Accept': 'application/vnd.ksql.v1+json'
    }
    
    if config.get('ksql_username') and config.get('ksql_password'):
        credentials = base64.b64encode(
            f"{config['ksql_username']}:{config['ksql_password']}".encode()
        ).decode()
        headers['Authorization'] = f'Basic {credentials}'
    
    return headers


@router.post("/query", response_model=KsqlQueryResponse)
async def execute_ksql_query(request: KsqlQueryRequest):
    """执行KSQL查询"""
    start_time = asyncio.get_event_loop().time()
    
    try:
        config = await get_ksql_config()
        headers = build_headers(config)
        
        # 构建请求体
        payload = {
            "ksql": request.ksql,
            "streamsProperties": request.streamsProperties or {}
        }
        
        if request.sessionVariables:
            payload["sessionVariables"] = request.sessionVariables
        
        timeout = aiohttp.ClientTimeout(total=config.get('ksql_timeout', 30000) / 1000)
        
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(
                f"{config['ksql_server_url']}/ksql",
                headers=headers,
                json=payload
            ) as response:
                execution_time = int((asyncio.get_event_loop().time() - start_time) * 1000)
                
                if response.status != 200:
                    error_text = await response.text()
                    return KsqlQueryResponse(
                        success=False,
                        error=f"HTTP {response.status}: {error_text}",
                        execution_time=execution_time
                    )
                
                result = await response.json()
                
                # 处理错误响应
                if result and len(result) > 0:
                    first_result = result[0]
                    
                    # 检查是否有错误
                    if 'errorMessage' in first_result:
                        error_msg = first_result['errorMessage']
                        if isinstance(error_msg, dict):
                            error_text = error_msg.get('message', str(error_msg))
                        else:
                            error_text = str(error_msg)
                        
                        return KsqlQueryResponse(
                            success=False,
                            error=error_text,
                            execution_time=execution_time
                        )
                    
                    # 处理查询结果
                    if 'queryResult' in first_result:
                        query_result = first_result['queryResult']
                        rows = query_result.get('rows', [])
                        
                        # 转换行数据为字典格式
                        if 'header' in query_result and 'schema' in query_result['header']:
                            columns = [col['name'] for col in query_result['header']['schema']['fields']]
                            formatted_rows = []
                            for row in rows:
                                if len(row) == len(columns):
                                    formatted_rows.append(dict(zip(columns, row)))
                                else:
                                    formatted_rows.append({'data': row})
                            
                            return KsqlQueryResponse(
                                success=True,
                                data=formatted_rows,
                                execution_time=execution_time,
                                row_count=len(formatted_rows),
                                query_id=query_result.get('queryId')
                            )
                        else:
                            return KsqlQueryResponse(
                                success=True,
                                data=[{'data': row} for row in rows],
                                execution_time=execution_time,
                                row_count=len(rows)
                            )
                    
                    # 处理命令状态
                    elif 'commandStatus' in first_result:
                        command_status = first_result['commandStatus']
                        return KsqlQueryResponse(
                            success=True,
                            data=[{
                                'status': command_status.get('status'),
                                'message': command_status.get('message', ''),
                                'statementText': command_status.get('statementText', '')
                            }],
                            execution_time=execution_time,
                            row_count=1,
                            command_status=command_status
                        )
                    
                    # 处理其他类型的响应
                    else:
                        return KsqlQueryResponse(
                            success=True,
                            data=result,
                            execution_time=execution_time,
                            row_count=len(result)
                        )
                
                # 空结果
                return KsqlQueryResponse(
                    success=True,
                    data=[],
                    execution_time=execution_time,
                    row_count=0
                )
                
    except asyncio.TimeoutError:
        execution_time = int((asyncio.get_event_loop().time() - start_time) * 1000)
        return KsqlQueryResponse(
            success=False,
            error="查询超时",
            execution_time=execution_time
        )
    except Exception as e:
        execution_time = int((asyncio.get_event_loop().time() - start_time) * 1000)
        logger.error(f"执行KSQL查询失败: {e}", exc_info=True)
        return KsqlQueryResponse(
            success=False,
            error=f"查询执行失败: {str(e)}",
            execution_time=execution_time
        )


@router.get("/info")
async def get_ksql_info():
    """获取KSQL服务器信息"""
    try:
        config = await get_ksql_config()
        headers = build_headers(config)

        timeout = aiohttp.ClientTimeout(total=10)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(f"{config['ksql_server_url']}/info", headers=headers) as response:
                if response.status == 200:
                    return await response.json()
                elif response.status == 404:
                    raise HTTPException(status_code=404, detail="KSQL服务器未找到，请检查服务器地址和端口")
                else:
                    error_text = await response.text()
                    # 简化HTML错误信息
                    if '<html>' in error_text.lower():
                        raise HTTPException(status_code=response.status, detail=f"KSQL服务器返回错误 (HTTP {response.status})")
                    else:
                        raise HTTPException(status_code=response.status, detail=error_text)

    except aiohttp.ClientConnectorError:
        logger.error("无法连接到KSQL服务器")
        raise HTTPException(status_code=503, detail="无法连接到KSQL服务器，请检查服务器是否运行")
    except asyncio.TimeoutError:
        logger.error("连接KSQL服务器超时")
        raise HTTPException(status_code=504, detail="连接KSQL服务器超时")
    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"获取KSQL服务器信息失败: {e}")
        raise HTTPException(status_code=500, detail="KSQL服务器连接失败")


@router.get("/healthcheck")
async def ksql_healthcheck():
    """KSQL健康检查"""
    try:
        config = await get_ksql_config()
        headers = build_headers(config)

        timeout = aiohttp.ClientTimeout(total=5)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(f"{config['ksql_server_url']}/healthcheck", headers=headers) as response:
                if response.status == 200:
                    result = await response.json()
                    return {
                        "healthy": result.get("isHealthy", False),
                        "details": result.get("details", {})
                    }
                else:
                    return {
                        "healthy": False,
                        "error": f"服务器返回错误 (HTTP {response.status})"
                    }

    except aiohttp.ClientConnectorError:
        return {
            "healthy": False,
            "error": "无法连接到KSQL服务器"
        }
    except asyncio.TimeoutError:
        return {
            "healthy": False,
            "error": "连接超时"
        }
    except Exception as e:
        logger.error(f"KSQL健康检查失败: {e}")
        return {
            "healthy": False,
            "error": "健康检查失败"
        }


@router.get("/status")
async def get_ksql_status():
    """获取KSQL集群状态"""
    try:
        config = await get_ksql_config()
        headers = build_headers(config)
        
        timeout = aiohttp.ClientTimeout(total=10)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.get(f"{config['ksql_server_url']}/status", headers=headers) as response:
                if response.status == 200:
                    return await response.json()
                else:
                    raise HTTPException(status_code=response.status, detail=await response.text())
                    
    except Exception as e:
        logger.error(f"获取KSQL状态失败: {e}")
        raise HTTPException(status_code=500, detail=f"获取状态失败: {str(e)}")


@router.post("/query/stream")
async def execute_stream_query(request: KsqlStreamRequest):
    """执行流查询（用于持续查询）"""
    try:
        config = await get_ksql_config()
        headers = build_headers(config)
        headers['Accept'] = 'application/vnd.ksql.v1+json'
        
        payload = {
            "sql": request.sql,
            "properties": request.properties or {}
        }
        
        timeout = aiohttp.ClientTimeout(total=config.get('ksql_timeout', 30000) / 1000)
        
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(
                f"{config['ksql_server_url']}/query-stream",
                headers=headers,
                json=payload
            ) as response:
                if response.status == 200:
                    # 对于流查询，我们返回查询ID和状态
                    return {
                        "success": True,
                        "message": "流查询已启动",
                        "query_id": f"stream_{asyncio.get_event_loop().time()}"
                    }
                else:
                    error_text = await response.text()
                    raise HTTPException(status_code=response.status, detail=error_text)
                    
    except Exception as e:
        logger.error(f"执行流查询失败: {e}")
        raise HTTPException(status_code=500, detail=f"流查询失败: {str(e)}")


@router.delete("/query/{query_id}")
async def terminate_query(query_id: str):
    """终止查询"""
    try:
        config = await get_ksql_config()
        headers = build_headers(config)
        
        payload = {
            "ksql": f"TERMINATE {query_id};",
            "streamsProperties": {}
        }
        
        timeout = aiohttp.ClientTimeout(total=10)
        async with aiohttp.ClientSession(timeout=timeout) as session:
            async with session.post(
                f"{config['ksql_server_url']}/ksql",
                headers=headers,
                json=payload
            ) as response:
                if response.status == 200:
                    return {"success": True, "message": f"查询 {query_id} 已终止"}
                else:
                    error_text = await response.text()
                    raise HTTPException(status_code=response.status, detail=error_text)
                    
    except Exception as e:
        logger.error(f"终止查询失败: {e}")
        raise HTTPException(status_code=500, detail=f"终止查询失败: {str(e)}")
